import numpy as np
import pickle
import cv2 as cv
import matplotlib.pyplot as plt

# 根据四元数计算位置
def LocationCalculate(X, Y, Z):
    a = np.arccos(np.sqrt(X ** 2 + Z ** 2) / np.sqrt(X ** 2 + Y ** 2 + Z ** 2))
    if Y > 0:
        ver = a / np.pi * 180
    else:
        ver = -a / np.pi * 180

    b = np.arccos(X / np.sqrt(X ** 2 + Z ** 2))
    if Z < 0:
        hor = b / np.pi * 180
    else:
        hor = (2. - b / np.pi) * 180

    return (90 - ver) / 180, hor / 360

data = pickle.load(open('./data/saliency_ds1_topicparis', 'rb'), encoding='bytes')

timestamp, fixation_list, saliency_map = data[0]
    
# print out the values of fields in the first record
print("timestamp:", type(timestamp), timestamp)
print("fixation_list:", type(fixation_list), len(fixation_list), 'fixation_list[0]:', fixation_list[0])
print("saliency_map:", type(saliency_map), saliency_map.shape)

""" cv.imshow('image[0]', saliency_map)
k = cv.waitKey(0) & 0xFF

if k == 27:                # wait for ESC key to exit
    cv.destroyAllWindows()
    print('Esc...')
elif k == ord('s'):        # wait for 's' key to save and exit
    cv.imwrite('./visualize/image[0].jpg', saliency_map)
    cv.destroyAllWindows()
    print('Save...') """

cap = cv.VideoCapture('/Users/pengshuai/Desktop/peng/360Degree_Head_Movement_Dataset/videos/Paris.mkv')
print(f'Paris FW={cap.get(3)} FH={cap.get(4)} FPS={cap.get(5)} Frames={cap.get(7)}')

for i, [timestamp, fixation_list, saliency_map] in enumerate(data, 0):
    if i < 100:
        continue

    cap.set(0, timestamp * 1000)

    _, frame = cap.read()
    if frame is None:
        print("Error for frame...")
    else:
        cv.imwrite(f'./visualize/frame[{i}].jpg', frame)
    

    points = []
    for [x, y, z] in fixation_list:
        points.append(list(LocationCalculate(x, y, z)))
    # print(len(fixation_list), len(points))
    points = np.array(points)
    # print(points.shape)


    plt.scatter(points[:, 0], points[:, 1])
    plt.xticks([]), plt.yticks([])    # to hide tick values on X and Y axis
    plt.savefig(f'./visualize/fix_image[{i}].jpg')
    # plt.show()
    plt.close()

    plt.imshow(saliency_map, cmap='gray')
    plt.xticks([]), plt.yticks([])    # to hide tick values on X and Y axis
    plt.savefig(f'./visualize/sal_image[{i}].jpg')
    # plt.show()
    plt.close()

    # print(f'i={i}, save done...')
    # if i == 20:
    #     break

print("done!")
cap.release()